package hbbft

import (
	"bufio"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"sync"
	"time"
)

// TCPNetwork is a real network implementation over TCP sockets.
// It keeps persistent outbound connections to peers and accepts inbound connections.
// Messages are length-prefixed (uint32 big endian) binary frames.
// No artificial delay/loss is introduced here.
type TCPNetwork struct {
	selfID    uint64
	selfAddr  string
	peerIDToAddr map[uint64]string

	ln       net.Listener
	mu       sync.RWMutex
	conns    map[uint64]net.Conn

	readTimeout  time.Duration
	writeTimeout time.Duration

	// onMessage is invoked when a message is received.
	// The first argument is the remote peer ID if known (0 if unknown),
	// and the second is the raw byte payload.
	onMessage func(from uint64, data []byte)
}

func NewTCPNetwork(selfID uint64, selfAddr string, peers map[uint64]string, onMessage func(from uint64, data []byte)) *TCPNetwork {
	return &TCPNetwork{
		selfID:    selfID,
		selfAddr:  selfAddr,
		peerIDToAddr: peers,
		conns:     make(map[uint64]net.Conn),
		readTimeout:  30 * time.Second,
		writeTimeout: 30 * time.Second,
		onMessage: onMessage,
	}
}

// Start begins listening and serving inbound connections.
func (n *TCPNetwork) Start() error {
	ln, err := net.Listen("tcp", n.selfAddr)
	if err != nil {
		return err
	}
	n.ln = ln
	go n.acceptLoop()
	return nil
}

func (n *TCPNetwork) acceptLoop() {
	for {
		conn, err := n.ln.Accept()
		if err != nil {
			if errors.Is(err, net.ErrClosed) {
				return
			}
			continue
		}
		go n.handleConn(conn)
	}
}

func (n *TCPNetwork) handleConn(conn net.Conn) {
	defer conn.Close()
	fmt.Printf("[TCP] Accepted connection from %s\n", conn.RemoteAddr())
	reader := bufio.NewReader(conn)
	for {
		if n.readTimeout > 0 {
			_ = conn.SetReadDeadline(time.Now().Add(n.readTimeout))
		}
		// Read frame length
		var length uint32
		if err := binary.Read(reader, binary.BigEndian, &length); err != nil {
			if err == io.EOF {
				fmt.Printf("[TCP] Connection closed from %s\n", conn.RemoteAddr())
				return
			}
			fmt.Printf("[TCP] Read error from %s: %v\n", conn.RemoteAddr(), err)
			return
		}
		if length == 0 || length > 64*1024*1024 {
			// Defensive limits
			fmt.Printf("[TCP] Invalid length %d from %s\n", length, conn.RemoteAddr())
			return
		}
		buf := make([]byte, length)
		if _, err := io.ReadFull(reader, buf); err != nil {
			fmt.Printf("[TCP] Read payload error from %s: %v\n", conn.RemoteAddr(), err)
			return
		}
		fmt.Printf("[TCP] Received %d bytes from %s\n", length, conn.RemoteAddr())
		// Unknown peer id for inbound; higher layer may embed sender id in payload if needed.
		if n.onMessage != nil {
			n.onMessage(0, buf)
		}
	}
}

// Send sends a message to the target peer id.
func (n *TCPNetwork) Send(to uint64, msg []byte) error {
	conn, err := n.getOrDial(to)
	if err != nil {
		return err
	}
	return writeFrame(conn, n.writeTimeout, msg)
}

// Broadcast sends the message to all known peers (excluding self).
func (n *TCPNetwork) Broadcast(msg []byte) error {
	n.mu.RLock()
	peers := make([]uint64, 0, len(n.peerIDToAddr))
	for id := range n.peerIDToAddr {
		if id == n.selfID { continue }
		peers = append(peers, id)
	}
	n.mu.RUnlock()
	var firstErr error
	for _, id := range peers {
		if err := n.Send(id, msg); err != nil && firstErr == nil {
			firstErr = err
		}
	}
	return firstErr
}

func (n *TCPNetwork) getOrDial(to uint64) (net.Conn, error) {
	n.mu.RLock()
	if c, ok := n.conns[to]; ok {
		n.mu.RUnlock()
		return c, nil
	}
	n.mu.RUnlock()
	addr, ok := n.peerIDToAddr[to]
	if !ok {
		return nil, errors.New("unknown peer id")
	}
	c, err := net.DialTimeout("tcp", addr, 10*time.Second)
	if err != nil {
		return nil, err
	}
	n.mu.Lock()
	n.conns[to] = c
	n.mu.Unlock()
	return c, nil
}

func writeFrame(conn net.Conn, timeout time.Duration, payload []byte) error {
	if timeout > 0 { _ = conn.SetWriteDeadline(time.Now().Add(timeout)) }
	var header [4]byte
	binary.BigEndian.PutUint32(header[:], uint32(len(payload)))
	if _, err := conn.Write(header[:]); err != nil { return err }
	if _, err := conn.Write(payload); err != nil { return err }
	return nil
}

// Close shuts down the listener and all outbound connections.
func (n *TCPNetwork) Close() error {
	if n.ln != nil { _ = n.ln.Close() }
	n.mu.Lock()
	for id, c := range n.conns {
		_ = c.Close()
		delete(n.conns, id)
	}
	n.mu.Unlock()
	return nil
}

